"""
File   : typeddict.py
Desc   : langchain 结构化输出: TypedDict
         注意: langchain-core>=0.2.26
Date   : 2024/12/03
Author : Tianyu Chen
https://python.langchain.com/docs/how_to/structured_output/#typeddict-or-json-schema
"""

from typing import Optional
from typing_extensions import Annotated, TypedDict
from langchain_openai import ChatOpenAI

# Model
llm = ChatOpenAI(model="gpt-4o-mini")

# TypedDict
class Joke(TypedDict):
    """Joke to tell user."""

    setup: Annotated[str, ..., "The setup of the joke"]

    # Alternatively, we could have specified setup as:

    # setup: str                    # no default, no description
    # setup: Annotated[str, ...]    # no default, no description
    # setup: Annotated[str, "foo"]  # default, no description

    punchline: Annotated[str, ..., "The punchline of the joke"]
    rating: Annotated[Optional[int], None, "How funny the joke is, from 1 to 10"]

# Runnable
structured_llm = llm.with_structured_output(Joke)

# invoke
a_joke = structured_llm.invoke("Tell me a joke about cats")
print(type(a_joke))  # <class 'dict'>
print(a_joke)  # {'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the mouse!', 'rating': 7}
print("setup:", a_joke["setup"])  # setup: Why was the cat sitting on the computer?
print("punchline:", a_joke["punchline"])  # punchline: Because it wanted to keep an eye on the mouse!
print("rating:", a_joke["rating"])  # rating: 7
